import cv2
import numpy as np
import time
import pygame
from moviepy.video.io.VideoFileClip import VideoFileClip
import sounddevice as sd
from threading import Thread

from electron_player.electron_player import ElectronPlayer
from electron_low_level.electron_low_level import ElectronLowLevel

class ElectronBotTest:
    def __init__(self):
        # 初始化底层控制
        self.robot = ElectronLowLevel()
        # 初始化播放器
        self.player = ElectronPlayer()
        
    def connect(self):
        """连接机器人"""
        if self.robot.connect():
            print("机器人连接成功!")
            return True
        else:
            print("机器人连接失败!")
            return False
            
    def test_video(self):
        """测试视频播放"""
        print("开始视频测试...")
        
        def play_audio(audio_clip):
            """在独立线程中播放音频"""
            try:
                # 获取音频数据
                audio_array = audio_clip.to_soundarray()
                sample_rate = audio_clip.fps
                
                # 确保音频数据是float32类型
                audio_array = audio_array.astype(np.float32)
                
                # 如果是单声道，转换为双声道
                if len(audio_array.shape) == 1:
                    audio_array = np.column_stack((audio_array, audio_array))
                
                # 播放音频
                sd.play(audio_array, sample_rate)
                sd.wait()  # 等待音频播放完成
                
            except Exception as e:
                print(f"音频播放错误: {e}")
        
        try:
            # 支持本地视频或网络视频URL
            # video_path = "test.mp4"
            video_path = "https://media.w3.org/2010/05/sintel/trailer.mp4"
            
            # 使用moviepy加载视频
            video_clip = VideoFileClip(video_path)
            
            # 如果视频有音频轨道，启动音频播放线程
            if video_clip.audio is not None:
                audio_thread = Thread(target=play_audio, args=(video_clip.audio,))
                audio_thread.daemon = True
                audio_thread.start()
            
            # 获取视频参数
            fps = video_clip.fps
            frame_interval = 1.0 / fps
            
            # 逐帧处理视频
            start_time = time.time()
            for frame in video_clip.iter_frames(fps=fps):
                current_time = time.time() - start_time
                
                # 发送帧到机器人显示
                self.robot.set_image_src(frame)
                self.robot.sync()
                
                # 计算下一帧的等待时间
                wait_time = frame_interval - (time.time() - start_time - current_time)
                if wait_time > 0:
                    time.sleep(wait_time)
                
                # 按'q'键退出
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
                
        finally:
            # 清理资源
            video_clip.close()
            sd.stop()
        
        print("视频测试完成")
        
    def test_joint_movement(self):
        """测试关节运动"""
        print("开始关节运动测试...")
        
        try:
            # 获取当前关节状态
            angles = [0.0] * 6
            self.robot.get_joint_angles(angles)
            print(f"当前关节角度: {angles}")
            
            # 测试设置关节角度 - 添加更长的延时和更大的角度
            test_angles = [120.0, 0.0, 0.0, 0.0, 0.0, 0.0]  # 增大角度到30度
            print("正在移动到测试位置...")
            self.robot.set_joint_angles(*test_angles, True)
            self.robot.sync()
            time.sleep(3)  # 增加等待时间到3秒
            
            # 读取新的角度
            self.robot.get_joint_angles(angles)
            print(f"移动后关节角度: {angles}")
            
            # 返回初始位置
            print("正在返回初始位置...")
            zero_angles = [0.0] * 6
            self.robot.set_joint_angles(*zero_angles, True)
            self.robot.sync()
            time.sleep(3)  # 增加等待时间到3秒
            
        except Exception as e:
            print(f"关节测试出错: {e}")
        
        print("关节运动测试完成")
        
    def test_audio_playback(self):
        """测试音频播放"""
        print("开始音频播放测试...")
        # 初始化pygame音频
        pygame.mixer.init()
        
        try:
            # 加载并播放音频
            pygame.mixer.music.load("test.mp3")
            pygame.mixer.music.play()
            
            # 等待音频播放完成
            while pygame.mixer.music.get_busy():
                time.sleep(0.1)
            
        finally:
            # 停止音频播放并清理资源
            pygame.mixer.music.stop()
            pygame.mixer.quit()
            print("音频播放测试完成")
        
    def test_camera(self):
        """测试摄像头"""
        print("开始摄像头测试...")
        cap = cv2.VideoCapture(0)
        
        for _ in range(100):  # 捕获100帧
            ret, frame = cap.read()
            if ret:
                self.robot.set_image_src(frame)
                self.robot.sync()
                time.sleep(0.033)
                
        cap.release()
        print("摄像头测试完成")
        
    def run_all_tests(self):
        """运行所有测试"""
        if not self.connect():
            return
            
        try:
            # self.test_joint_movement()
            # self.test_video()
            self.test_audio_playback()
            # self.test_camera()
            
        finally:
            self.robot.disconnect()
            print("测试完成，机器人已断开连接")

if __name__ == "__main__":
    test = ElectronBotTest()
    test.run_all_tests() 